# def dataset_kwargs(cfg):
#     """
#     Build kwargs for TrainImageDataManager in image_data_manager.py from
#     the parsed command-line arguments.
#     """
#
#     return {
#         'dataset_name': cfg.INPUT.DATASET_NAME,
#         'root': cfg.params.TRAIN_IMAGES_PATH,
#         'height': cfg.INPUT.HEIGHT,
#         'width': cfg.INPUT.WIDTH,
#         'train_batch_size': cfg.params.BATCH_SIZE,
#         'test_batch_size': cfg.params.BATCH_SIZE,
#         'workers': cfg.params.WORKERS,
#         'train_sampler': cfg.TRAIN_SAMPLER,
#         'random_erase': cfg.DATAAUG.RANDOM_ERASE,
#         'color_jitter': cfg.DATAAUG.COLOR_JITTER,
#         'color_aug': cfg.DATAAUG.COLOR_AUG,
#     }


def optimizer_kwargs(cfg):
    """
    Build kwargs for optimizer in optimizers.py from
    the parsed command-line arguments.
    """
    return {
        'optim': cfg.OPTIM.FEM_METHOD,
        'lr': cfg.OPTIM.LR,
        'weight_decay': cfg.OPTIM.WEIGHT_DECAY,
        'momentum': cfg.OPTIM.SGD.MOMENTUM,
        'sgd_dampening': cfg.OPTIM.SGD.DAMPEN,
        'sgd_nesterov': cfg.OPTIM.SGD.NESTEROV,
        'rmsprop_alpha': cfg.OPTIM.RMSPROP,
        'adam_beta1': cfg.OPTIM.ADAM.BETA1,
        'adam_beta2': cfg.OPTIM.ADAM.BETA2,
        'staged_lr': cfg.OPTIM.STAGED_LR,
    }


def lr_scheduler_kwargs(cfg):
    """
    Build kwargs for lr_scheduler in lr_schedulers.py from
    the parsed command-line arguments.
    """
    return {
        'lr_scheduler': cfg.OPTIM.LRSCHEDULER.METHOD,
        'stepsize': cfg.OPTIM.LRSCHEDULER.STEPSIZE,
        'gamma': cfg.OPTIM.LRSCHEDULER.GAMMA,
        'warmup_iters': cfg.OPTIM.LRSCHEDULER.WARMUP_ITERS,
        'warmup_factor': cfg.OPTIM.LRSCHEDULER.WARMUP_FACTOR,
        'warmup_method': cfg.OPTIM.LRSCHEDULER.WARMUP_METHOD,
    }
